#!/usr/bin/env python
#  -*- coding: utf-8 -*-
"""
包含有关时间处理的功能函数
"""

__author__ = 'limin'

import datetime
import time

import pandas as pd


# 2003-01-01 ～ 2021-12-31 节假日，不包含周末
chinese_rest_days = ['2003-01-01', '2003-01-30', '2003-01-31', '2003-02-03', '2003-02-04', '2003-02-05', '2003-02-06', '2003-02-07', '2003-05-01', '2003-05-02', '2003-05-05', '2003-05-06', '2003-05-07', '2003-10-01', '2003-10-02', '2003-10-03', '2003-10-06', '2003-10-07', '2004-01-01', '2004-01-19', '2004-01-20', '2004-01-21', '2004-01-22', '2004-01-23', '2004-01-26', '2004-01-27', '2004-01-28', '2004-05-03', '2004-05-04', '2004-05-05', '2004-05-06', '2004-05-07', '2004-10-01', '2004-10-04', '2004-10-05', '2004-10-06', '2004-10-07', '2005-01-03', '2005-02-07', '2005-02-08', '2005-02-09', '2005-02-10', '2005-02-11', '2005-02-14', '2005-02-15', '2005-05-02', '2005-05-03', '2005-05-04', '2005-05-05', '2005-05-06', '2005-10-03', '2005-10-04', '2005-10-05', '2005-10-06', '2005-10-07', '2006-01-02', '2006-01-03', '2006-01-30', '2006-01-31', '2006-02-01', '2006-02-02', '2006-02-03', '2006-05-01', '2006-05-02', '2006-05-03', '2006-05-04', '2006-05-05', '2006-10-02', '2006-10-03', '2006-10-04', '2006-10-05', '2006-10-06', '2007-01-01', '2007-01-02', '2007-01-03', '2007-02-19', '2007-02-20', '2007-02-21', '2007-02-22', '2007-02-23', '2007-05-01', '2007-05-02', '2007-05-03', '2007-05-04', '2007-05-07', '2007-10-01', '2007-10-02', '2007-10-03', '2007-10-04', '2007-10-05', '2007-12-31', '2008-01-01', '2008-02-06', '2008-02-07', '2008-02-08', '2008-02-11', '2008-02-12', '2008-04-04', '2008-05-01', '2008-05-02', '2008-06-09', '2008-09-15', '2008-09-29', '2008-09-30', '2008-10-01', '2008-10-02', '2008-10-03', '2009-01-01', '2009-01-02', '2009-01-26', '2009-01-27', '2009-01-28', '2009-01-29', '2009-01-30', '2009-04-06', '2009-05-01', '2009-05-28', '2009-05-29', '2009-10-01', '2009-10-02', '2009-10-05', '2009-10-06', '2009-10-07', '2009-10-08', '2010-01-01', '2010-02-15', '2010-02-16', '2010-02-17', '2010-02-18', '2010-02-19', '2010-04-05', '2010-05-03', '2010-06-14', '2010-06-15', '2010-06-16', '2010-09-22', '2010-09-23', '2010-09-24', '2010-10-01', '2010-10-04', '2010-10-05', '2010-10-06', '2010-10-07', '2011-01-03', '2011-02-02', '2011-02-03', '2011-02-04', '2011-02-07', '2011-02-08', '2011-04-04', '2011-04-05', '2011-05-02', '2011-06-06', '2011-09-12', '2011-10-03', '2011-10-04', '2011-10-05', '2011-10-06', '2011-10-07', '2012-01-02', '2012-01-03', '2012-01-23', '2012-01-24', '2012-01-25', '2012-01-26', '2012-01-27', '2012-04-02', '2012-04-03', '2012-04-04', '2012-04-30', '2012-05-01', '2012-06-22', '2012-10-01', '2012-10-02', '2012-10-03', '2012-10-04', '2012-10-05', '2013-01-01', '2013-01-02', '2013-01-03', '2013-02-11', '2013-02-12', '2013-02-13', '2013-02-14', '2013-02-15', '2013-04-04', '2013-04-05', '2013-04-29', '2013-04-30', '2013-05-01', '2013-06-10', '2013-06-11', '2013-06-12', '2013-09-19', '2013-09-20', '2013-10-01', '2013-10-02', '2013-10-03', '2013-10-04', '2013-10-07', '2014-01-01', '2014-01-31', '2014-02-03', '2014-02-04', '2014-02-05', '2014-02-06', '2014-04-07', '2014-05-01', '2014-05-02', '2014-06-02', '2014-09-08', '2014-10-01', '2014-10-02', '2014-10-03', '2014-10-06', '2014-10-07', '2015-01-01', '2015-01-02', '2015-02-18', '2015-02-19', '2015-02-20', '2015-02-23', '2015-02-24', '2015-04-06', '2015-05-01', '2015-06-22', '2015-09-03', '2015-09-04', '2015-10-01', '2015-10-02', '2015-10-05', '2015-10-06', '2015-10-07', '2016-01-01', '2016-02-08', '2016-02-09', '2016-02-10', '2016-02-11', '2016-02-12', '2016-04-04', '2016-05-02', '2016-06-09', '2016-06-10', '2016-09-15', '2016-09-16', '2016-10-03', '2016-10-04', '2016-10-05', '2016-10-06', '2016-10-07', '2017-01-02', '2017-01-27', '2017-01-30', '2017-01-31', '2017-02-01', '2017-02-02', '2017-04-03', '2017-04-04', '2017-05-01', '2017-05-29', '2017-05-30', '2017-10-02', '2017-10-03', '2017-10-04', '2017-10-05', '2017-10-06', '2018-01-01', '2018-02-15', '2018-02-16', '2018-02-19', '2018-02-20', '2018-02-21', '2018-04-05', '2018-04-06', '2018-04-30', '2018-05-01', '2018-06-18', '2018-09-24', '2018-10-01', '2018-10-02', '2018-10-03', '2018-10-04', '2018-10-05', '2018-12-31', '2019-01-01', '2019-02-04', '2019-02-05', '2019-02-06', '2019-02-07', '2019-02-08', '2019-04-05', '2019-05-01', '2019-05-02', '2019-05-03', '2019-06-07', '2019-09-13', '2019-10-01', '2019-10-02', '2019-10-03', '2019-10-04', '2019-10-07', '2020-01-01', '2020-01-24', '2020-01-27', '2020-01-28', '2020-01-29', '2020-01-30', '2020-01-31', '2020-04-06', '2020-05-01', '2020-05-04', '2020-05-05', '2020-06-25', '2020-06-26', '2020-10-01', '2020-10-02', '2020-10-05', '2020-10-06', '2020-10-07', '2020-10-08', '2021-01-01', '2021-02-11', '2021-02-12', '2021-02-15', '2021-02-16', '2021-02-17', '2021-04-05', '2021-05-03', '2021-05-04', '2021-05-05', '2021-06-14', '2021-09-20', '2021-09-21', '2021-10-01', '2021-10-04', '2021-10-05', '2021-10-06', '2021-10-07']
rest_days_df = pd.DataFrame()
rest_days_df['date'] = pd.Series(pd.to_datetime(chinese_rest_days, format='%Y-%m-%d'))
rest_days_df['trading'] = False


def _get_trading_calendar(start_dt, end_dt):
    """获取一段时间内，每天是否是交易日"""
    df = pd.DataFrame()
    df['date'] = pd.Series(pd.date_range(start=start_dt, end=end_dt, freq="D"))
    df['trading'] = df['date'].dt.dayofweek.lt(5)
    result = pd.merge(rest_days_df, df, sort=True, how="right", on="date")
    result.fillna(True, inplace=True)
    df['trading'] = result['trading_x'] & result['trading_y']
    return df


def _format_from_timestamp_nano(nano, fmt="%Y-%m-%d %H:%M:%S.%f"):
    return datetime.datetime.fromtimestamp(nano / 1e9).strftime(fmt)


def _get_trading_day_start_time(trading_day):
    """给定交易日, 获得交易日起始时间"""
    begin_mark = 631123200000000000  # 1990-01-01
    start_time = trading_day - 21600000000000  # 6小时
    week_day = (start_time - begin_mark) // 86400000000000 % 7
    if week_day >= 5:
        start_time -= 86400000000000 * (week_day - 4)
    return start_time


def _get_trading_day_end_time(trading_day):
    """给定交易日, 获得交易日结束时间"""
    return trading_day + 64799999999999  # 18小时


def _get_trading_day_from_timestamp(timestamp):
    """给定时间, 获得其所属的交易日"""
    begin_mark = 631123200000000000  # 1990-01-01
    days = (timestamp - begin_mark) // 86400000000000  # 自 1990-01-01 以来的天数
    if (timestamp - begin_mark) % 86400000000000 >= 64800000000000:  # 大于18点, 天数+1
        days += 1
    week_day = days % 7
    if week_day >= 5:  # 如果是周末则移到星期一
        days += 7 - week_day
    return begin_mark + days * 86400000000000


def _get_trading_timestamp(quote, current_datetime: str):
    """ 将 quote 在 current_datetime 所在交易日的所有可交易时间段转换为纳秒时间戳(tqsdk内部使用的时间戳统一为纳秒)并返回 """
    # 获取当前交易日时间戳
    current_trading_day_timestamp = _get_trading_day_from_timestamp(
        int(datetime.datetime.strptime(current_datetime, "%Y-%m-%d %H:%M:%S.%f").timestamp() * 1e6) * 1000)
    # 获取上一交易日时间戳
    last_trading_day_timestamp = _get_trading_day_from_timestamp(
        _get_trading_day_start_time(current_trading_day_timestamp) - 1)
    trading_timestamp = {
        "day": _get_period_timestamp(current_trading_day_timestamp, quote["trading_time"].get("day", [])),
        "night": _get_period_timestamp(last_trading_day_timestamp, quote["trading_time"].get("night", []))
    }
    return trading_timestamp


def _get_period_timestamp(real_date_timestamp, period_str):
    """
    real_date_timestamp：period_str 所在真实日期的纳秒时间戳（如 period_str 为周一(周二)的夜盘,则real_date_timestamp为上周五(周一)的日期; period_str 为周一的白盘,则real_date_timestamp为周一的日期）
    period_str: quote["trading_time"]["day"] or quote["trading_time"]["night"]
    """
    period_timestamp = []
    for duration in period_str:  # 对于白盘（或夜盘）中的每一个可交易时间段
        start = [int(i) for i in duration[0].split(":")]  # 交易时间段起始点
        end = [int(i) for i in duration[1].split(":")]  # 交易时间段结束点
        period_timestamp.append([real_date_timestamp + (start[0] * 3600 + start[1] * 60 + start[2]) * 1000000000,
                                 real_date_timestamp + (end[0] * 3600 + end[1] * 60 + end[2]) * 1000000000])
    return period_timestamp


def _is_in_trading_time(quote, current_datetime, local_time_record):
    """ 判断是否在可交易时间段内，需在quote已收到行情后调用本函数"""
    # 只在需要用到可交易时间段时(即本函数中)才调用_get_trading_timestamp()
    trading_timestamp = _get_trading_timestamp(quote, current_datetime)
    now_ns_timestamp = _get_trade_timestamp(current_datetime, local_time_record)  # 当前预估交易所纳秒时间戳
    # 判断当前交易所时间（估计值）是否在交易时间段内
    for v in trading_timestamp.values():
        for period in v:
            if period[0] <= now_ns_timestamp < period[1]:
                return True
    return False


def _get_trade_timestamp(current_datetime, local_time_record):
    # 根据最新行情时间获取模拟的(预估的)当前交易所纳秒时间戳（tqsdk内部使用的时间戳统一为纳秒）
    # 如果local_time_record为nan，則不加时间差
    return int((datetime.datetime.strptime(current_datetime,
                                           "%Y-%m-%d %H:%M:%S.%f").timestamp() + (
                    0 if local_time_record != local_time_record else (
                            time.time() - local_time_record))) * 1e6) * 1000


def _get_expire_rest_days(expire_dt, current_dt):
    """
    获取当前时间到下市时间之间的天数
    expire_dt, current_dt 都以 s 为单位
    """
    delta = datetime.datetime.fromtimestamp(expire_dt).date() - datetime.datetime.fromtimestamp(current_dt).date()
    return delta.days
